from langchain.agents import create_react_agent, AgentExecutor
from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain_community.document_loaders import JSONLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_chroma import Chroma
from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from langchain_text_splitters import CharacterTextSplitter
import os

# 定义数据库目录
persist_directory = "./chroma_db"

# 检查数据库是否已经存在
if not os.path.exists(persist_directory):
    # 加载数据集
    doc = JSONLoader(
        file_path="data.json",
        jq_schema=".[].content+.[].answer",
        text_content=True,
    ).load()  # 加载数据集
    print("数据集加载完成", doc)
    # 创建文本切割器
    splitter = CharacterTextSplitter('\n', chunk_size=100, chunk_overlap=0)  # 设置切割大小和重叠大小

    # 切割数据集
    chunks = splitter.split_documents(doc)  # 将加载的数据集切割成小块

    # 创建文本嵌入模型
    embeddings = DashScopeEmbeddings()  # 使用 DashScopeEmbeddings，这是一个开源的文本嵌入模型，通过调用 DashScope 的文本嵌入服务，将文本转换为向量表示

    # 将切割后的数据集转换为向量表示并存储到本地数据库中
    db = Chroma.from_documents(chunks, embeddings, persist_directory=persist_directory)
else:
    # 从本地数据库中加载数据
    embeddings = DashScopeEmbeddings()
    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
# 创建 RetrievalQA 实例
qa = RetrievalQA.from_chain_type(llm=Tongyi(), chain_type='stuff', retriever=db.as_retriever())


# # 生成答案
# print(qa.invoke('小米15售价多少'))



# def get_answer(question):
#     # 生成答案
#     res = db.similarity_search(question, k=3)
#     return ''
#
#
# tool1 = Tool(
#     func=get_answer,
#     name='get_answer',
#     description='获取问题的答案',
# )
# tools = [tool1, ]
#
# template = '''请尽可能回答以下问题。您可以使用以下工具:
#
#             {tools}
#
#            使用以下格式:
#
#             Question: the input question you must answer
#             Thought: you should always think about what to do
#             Action: the action to take, should be one of [{tool_names}]
#             Action Input: the input to the action
#             Observation: the result of the action
#             ... (this Thought/Action/Action Input/Observation can repeat N times)
#             Thought: I now know the final answer
#             Final Answer: the final answer to the original input question
#
#             Begin!
#
#             Question: {input}
#             Thought:{agent_scratchpad}'''
#
# prompt = ChatPromptTemplate.from_template(template)
# agent = create_react_agent(llm=Tongyi(), tools=tools, prompt=prompt)
# agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
#
# res = agent_executor.invoke(
#     {'input': 'NBA冠军队是那一支球队'},
#     config={"configurer": {"session_id": "test"}}
# )
# print(res)
